{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DPL 模式重写"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T\n",
    "from tvm import relax as rx\n",
    "from tvm import relay, tir\n",
    "from tvm.relax.analysis import get_var2val\n",
    "from tvm.relax.dpl import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 简单的示例"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "原始表达式为：\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "&x2 = x + x\\\\\n",
    "&x4 = x2 + x2\n",
    "\\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "@R.function\n",
    "def main(x: R.Tensor((16, 16), \"float32\")) -> R.Tensor((16, 16), \"float32\"):\n",
    "    with R.dataflow():\n",
    "        x2 = R.add(x, x)\n",
    "        x4 = R.add(x2, x2)\n",
    "        R.output(x4)\n",
    "    return x4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "    <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "        x2: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(x, x)\n",
       "        x4: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(x2, x2)\n",
       "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(x4)\n",
       "    <span style=\"color: #008000; font-weight: bold\">return</span> x4\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "main.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "构建模板："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = wildcard()\n",
    "pattern = is_op(\"relax.add\")(x, x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "重写模式："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rewriter(_, matchings):\n",
    "    return R.multiply(matchings[x], R.const(2, \"float32\"))\n",
    "rewritten = rewrite_call(pattern, rewriter, main)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "    <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "        x2: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>multiply(x, R<span style=\"color: #AA22FF; font-weight: bold\">.</span>const(<span style=\"color: #008000\">2.0</span>, <span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        x4: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>multiply(x2, R<span style=\"color: #AA22FF; font-weight: bold\">.</span>const(<span style=\"color: #008000\">2.0</span>, <span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(x4)\n",
       "    <span style=\"color: #008000; font-weight: bold\">return</span> x4\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "rewritten.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "此时将 $x + x$ 重写为 $x \\times 2$ 。更进一步，原式可以化简为："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "    <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "        x4: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>multiply(x, R<span style=\"color: #AA22FF; font-weight: bold\">.</span>const(<span style=\"color: #008000\">4.0</span>, <span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(x4)\n",
       "    <span style=\"color: #008000; font-weight: bold\">return</span> x4\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "add1 = is_op(\"relax.add\")(x, x)\n",
    "pattern = is_op(\"relax.add\")(add1, add1)\n",
    "\n",
    "def rewriter(_, matchings):\n",
    "    return R.multiply(matchings[x], R.const(4, \"float32\"))\n",
    "\n",
    "rewritten = rewrite_call(pattern, rewriter, main)\n",
    "rewritten.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "不进行重写，按原样返回原始调用节点："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rewriter(orig, _):\n",
    "    return orig\n",
    "\n",
    "rewritten = rewrite_call(pattern, rewriter, main)\n",
    "tvm.ir.assert_structural_equal(rewritten, main)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 重写注意力模块"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(Q: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), K: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), V: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "    <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "        lv58: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(Q, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>])\n",
       "        lv59: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>reshape(lv58, R<span style=\"color: #AA22FF; font-weight: bold\">.</span>shape([<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">40</span>]))\n",
       "        lv61: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(K, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>])\n",
       "        lv62: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>reshape(lv61, R<span style=\"color: #AA22FF; font-weight: bold\">.</span>shape([<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">40</span>]))\n",
       "        lv64: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(V, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>])\n",
       "        lv65: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>reshape(lv64, R<span style=\"color: #AA22FF; font-weight: bold\">.</span>shape([<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">40</span>]))\n",
       "        lv62_transposed: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">40</span>, <span style=\"color: #008000\">4096</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(lv62, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>])\n",
       "        lv3_1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">4096</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(lv59, lv62_transposed, out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
       "        lv68: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">4096</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>multiply(lv3_1, R<span style=\"color: #AA22FF; font-weight: bold\">.</span>const(<span style=\"color: #008000\">0.15811388194561005</span>, <span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        lv69: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">4096</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>softmax(lv68, axis<span style=\"color: #AA22FF; font-weight: bold\">=-</span><span style=\"color: #008000\">1</span>)\n",
       "        lv_3: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(lv69, lv65, out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
       "        lv71: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>reshape(lv_3, R<span style=\"color: #AA22FF; font-weight: bold\">.</span>shape([<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">40</span>]))\n",
       "        lv72: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(lv71, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>])\n",
       "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(lv72)\n",
       "    <span style=\"color: #008000; font-weight: bold\">return</span> lv72\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@R.function\n",
    "def main(\n",
    "    Q: R.Tensor((2, 4096, 8, 40), \"float32\"),\n",
    "    K: R.Tensor((2, 4096, 8, 40), \"float32\"),\n",
    "    V: R.Tensor((2, 4096, 8, 40), \"float32\"),\n",
    ") -> R.Tensor((2, 4096, 8, 40), \"float32\"):\n",
    "    with R.dataflow():\n",
    "        lv58 = R.permute_dims(Q, axes=[0, 2, 1, 3])\n",
    "        lv59 = R.reshape(lv58, R.shape([16, 4096, 40]))\n",
    "\n",
    "        lv61 = R.permute_dims(K, axes=[0, 2, 1, 3])\n",
    "        lv62 = R.reshape(lv61, R.shape([16, 4096, 40]))\n",
    "\n",
    "        lv64 = R.permute_dims(V, axes=[0, 2, 1, 3])\n",
    "        lv65 = R.reshape(lv64, R.shape([16, 4096, 40]))\n",
    "\n",
    "        lv62_transposed = R.permute_dims(lv62, axes=[0, 2, 1])\n",
    "        lv3_1 = R.matmul(lv59, lv62_transposed)\n",
    "        lv68 = R.multiply(lv3_1, R.const(0.15811388194561005, \"float32\"))\n",
    "        lv69 = R.nn.softmax(lv68, axis=-1)\n",
    "        lv_3 = R.matmul(lv69, lv65)\n",
    "\n",
    "        lv71 = R.reshape(lv_3, R.shape([2, 8, 4096, 40]))\n",
    "        lv72 = R.permute_dims(lv71, axes=[0, 2, 1, 3])\n",
    "        R.output(lv72)\n",
    "\n",
    "    return lv72\n",
    "main.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "构建模板："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def BSNH_to_BSH(tensor):\n",
    "    return is_op(\"relax.reshape\")(is_op(\"relax.permute_dims\")(tensor), wildcard())\n",
    "\n",
    "def BSH_to_BSNH(tensor):\n",
    "    return is_op(\"relax.permute_dims\")(is_op(\"relax.reshape\")(tensor, wildcard()))\n",
    "\n",
    "Q = wildcard()\n",
    "K = wildcard()\n",
    "V = wildcard()\n",
    "\n",
    "Q_3D = BSNH_to_BSH(Q)\n",
    "V_3D = BSNH_to_BSH(V)\n",
    "K_3D = BSNH_to_BSH(K)\n",
    "\n",
    "matmul1 = is_op(\"relax.matmul\")(Q_3D, is_op(\"relax.permute_dims\")(V_3D))\n",
    "multiply = is_op(\"relax.multiply\")(matmul1, is_const())\n",
    "softmax = is_op(\"relax.nn.softmax\")(multiply)\n",
    "matmul2 = is_op(\"relax.matmul\")(softmax, K_3D)\n",
    "pattern = BSH_to_BSNH(matmul2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(Q: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), K: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), V: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "    <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "        lv72: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4096</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">40</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>attention(Q, V, K, scale<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>, causal_mask<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>, window_size<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
       "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(lv72)\n",
       "    <span style=\"color: #008000; font-weight: bold\">return</span> lv72\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def rewriter(_, matchings):\n",
    "    return R.nn.attention(matchings[Q], matchings[K], matchings[V])\n",
    "\n",
    "rewritten = rewrite_call(pattern, rewriter, main)\n",
    "rewritten.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试交换律模式匹配"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,))) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)):\n",
       "    <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "        y: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(x, x)\n",
       "        out: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(R<span style=\"color: #AA22FF; font-weight: bold\">.</span>const(<span style=\"color: #008000\">1.0</span>, <span style=\"color: #BA2121\">&quot;float32&quot;</span>), y)\n",
       "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(out)\n",
       "    <span style=\"color: #008000; font-weight: bold\">return</span> out\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@R.function(private=True)\n",
    "def before(\n",
    "    x: R.Tensor((1024,)),\n",
    "):\n",
    "    with R.dataflow():\n",
    "        y = R.add(x, x)\n",
    "        out = R.add(R.const(1.0), y)\n",
    "        R.output(out)\n",
    "    return out\n",
    "before.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,))) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)):\n",
       "    <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "        y: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(x, x)\n",
       "        out: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(y, R<span style=\"color: #AA22FF; font-weight: bold\">.</span>const(<span style=\"color: #008000\">2.0</span>, <span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(out)\n",
       "    <span style=\"color: #008000; font-weight: bold\">return</span> out\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pattern_add = is_op(\"relax.add\")\n",
    "pattern_mul = is_op(\"relax.multiply\")\n",
    "pattern_op = pattern_add | pattern_mul\n",
    "pattern_arg = wildcard()\n",
    "pattern_const = is_const()\n",
    "\n",
    "pattern = pattern_op(pattern_arg, pattern_const)\n",
    "\n",
    "def rewriter(expr, matches):\n",
    "    op = matches[pattern_op]\n",
    "    arg = matches[pattern_arg]\n",
    "    const = matches[pattern_const].data.numpy()\n",
    "    if const.shape == tuple() and const[()] == 1.0:\n",
    "        return rx.Call(op, [arg, rx.const(2.0)])\n",
    "    else:\n",
    "        return expr\n",
    "\n",
    "after = rewrite_call(pattern, rewriter, before)\n",
    "after.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试重复模式匹配"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "重写调用应迭代直到收敛:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)), y: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)), z: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,))) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)):\n",
       "    <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "        a: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(x, y)\n",
       "        b: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(a, z)\n",
       "        out: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>multiply(b, R<span style=\"color: #AA22FF; font-weight: bold\">.</span>const(<span style=\"color: #008000\">5.0</span>, <span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(out)\n",
       "    <span style=\"color: #008000; font-weight: bold\">return</span> out\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@R.function(private=True)\n",
    "def before(\n",
    "    x: R.Tensor((1024,)),\n",
    "    y: R.Tensor((1024,)),\n",
    "    z: R.Tensor((1024,)),\n",
    "):\n",
    "    with R.dataflow():\n",
    "        a = R.add(x, y)\n",
    "        b = R.add(a, z)\n",
    "        out = R.multiply(b, R.const(5.0))\n",
    "        R.output(out)\n",
    "    return out\n",
    "before.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)), y: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)), z: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,))) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)):\n",
       "    <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "        lv3: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>multiply(x, R<span style=\"color: #AA22FF; font-weight: bold\">.</span>const(<span style=\"color: #008000\">5.0</span>, <span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        lv4: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>multiply(y, R<span style=\"color: #AA22FF; font-weight: bold\">.</span>const(<span style=\"color: #008000\">5.0</span>, <span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        lv1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(lv3, lv4)\n",
       "        lv2: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>multiply(z, R<span style=\"color: #AA22FF; font-weight: bold\">.</span>const(<span style=\"color: #008000\">5.0</span>, <span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        out: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1024</span>,)) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(lv1, lv2)\n",
       "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(out)\n",
       "    <span style=\"color: #008000; font-weight: bold\">return</span> out\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pattern_add_lhs = wildcard()\n",
    "pattern_add_rhs = wildcard()\n",
    "pattern_add = is_op(\"relax.add\")(pattern_add_lhs, pattern_add_rhs)\n",
    "\n",
    "mul_const = is_const()\n",
    "pattern_mul = is_op(\"relax.multiply\")(pattern_add, mul_const)\n",
    "\n",
    "pattern = pattern_mul\n",
    "\n",
    "def rewriter(_expr, matches):\n",
    "    const = matches[mul_const]\n",
    "    return (matches[pattern_add_lhs] * const) + (matches[pattern_add_rhs] * const)\n",
    "\n",
    "after = rewrite_call(pattern, rewriter, before)\n",
    "after.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
